#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <getopt.h>

#define BUFFER_SIZE 2048

// 解析IP和端口
int parse_ip_port(const char *str, char *ip, int *port) {
    char *colon = strchr(str, ':');
    if (!colon) {
        fprintf(stderr, "Invalid format: %s\n", str);
        return -1;
    }
    *colon = '\0'; // 将字符串分割为IP和端口部分
    strcpy(ip, str);
    *port = atoi(colon + 1);
    return 0;
}

void print_usage() {
    printf("Usage: udp_forward -i <listen_ip:port> -o <target_ip:port> \n");
    printf("       udp_forward --input=<listen_ip:port> --output=<target_ip:port> \n");   
}

int main(int argc, char *argv[]) {
    char listen_ip[INET_ADDRSTRLEN] = "0.0.0.0";
    char target_ip[INET_ADDRSTRLEN] = "127.0.0.1";
    int listen_port = 5601, target_port = 5600;

    // 定义长选项
    struct option long_options[] = {
        {"input", required_argument, NULL, 'i'},
        {"output", required_argument, NULL, 'o'},
        {NULL, 0, NULL, 0}
    };

    int option_index = 0;
    int c;
    while ((c = getopt_long(argc, argv, "i:o:", long_options, &option_index)) != -1) {
        switch (c) {
            case 'i':
                if (parse_ip_port(optarg, listen_ip, &listen_port) != 0) {
                    print_usage();
                    return EXIT_FAILURE;
                }
                break;
            case 'o':
                if (parse_ip_port(optarg, target_ip, &target_port) != 0) {
                    print_usage();
                    return EXIT_FAILURE;
                }
                break;
            case '?':
                fprintf(stderr, "Unknown option: %c\n", optopt);
                print_usage();
                return EXIT_FAILURE;
            default:
                print_usage();
                return EXIT_FAILURE;
        }
    }

    // 检查是否提供了所有必要的参数
    if (listen_port == 0 || target_port == 0) {
        fprintf(stderr, "Usage: %s -i <listen_ip:port> -o <target_ip:port>\n", argv[0]);
        fprintf(stderr, "       %s --input=<listen_ip:port> --output=<target_ip:port>\n", argv[0]);
        return EXIT_FAILURE;
    }

    // 初始化监听地址
    struct sockaddr_in listen_addr, target_addr;
    memset(&listen_addr, 0, sizeof(listen_addr));
    listen_addr.sin_family = AF_INET;
    listen_addr.sin_port = htons(listen_port);
    if (inet_pton(AF_INET, listen_ip, &listen_addr.sin_addr) <= 0) {
        fprintf(stderr, "Invalid listen IP address: %s\n", listen_ip);
        return EXIT_FAILURE;
    }

    // 创建监听UDP套接字
    int listen_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (listen_fd < 0) {
        perror("Failed to create listen socket");
        return EXIT_FAILURE;
    }

    // 绑定监听地址
    if (bind(listen_fd, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) < 0) {
        perror("Failed to bind listen socket");
        close(listen_fd);
        return EXIT_FAILURE;
    }
    printf("Listening on %s:%d\n", listen_ip, listen_port);

    // 初始化目标地址
    memset(&target_addr, 0, sizeof(target_addr));
    target_addr.sin_family = AF_INET;
    target_addr.sin_port = htons(target_port);
    if (inet_pton(AF_INET, target_ip, &target_addr.sin_addr) <= 0) {
        fprintf(stderr, "Invalid target IP address: %s\n", target_ip);
        close(listen_fd);
        return EXIT_FAILURE;
    }
    printf("Sending to %s:%d\n", target_ip, target_port);

    char buffer[BUFFER_SIZE];
    socklen_t addr_len = sizeof(listen_addr);
    ssize_t bytes_received;
    // 主循环
    while (1) {
        // 接收数据
        bytes_received = recvfrom(listen_fd, buffer, BUFFER_SIZE, 0, 
                                          (struct sockaddr *)&listen_addr, &addr_len);
        if (bytes_received < 0) {
            perror("Failed to receive data");
            continue;
        }

        // 打印接收到的数据（可选）
        // printf("Received %ld bytes from %s:%d\n", bytes_received, 
        //       inet_ntoa(listen_addr.sin_addr), ntohs(listen_addr.sin_port));

        // 转发数据
        if (sendto(listen_fd, buffer, bytes_received, 0, 
                   (struct sockaddr *)&target_addr, sizeof(target_addr)) < 0) {
            perror("Failed to send data to target");
        } 
        /*else {
            printf("Forwarded %ld bytes to %s:%d\n", bytes_received, target_ip, target_port);
        }*/
    }

    // 关闭套接字
    close(listen_fd);
    return EXIT_SUCCESS;
}